#include "defines.h"
#include "lib.h"
#include "kozos.h"
#include "netdef.h"
#include "netlib.h"
#include "ni.h"
#include "icmp.h"
#include "udp.h"

static const uint8 loopback_address[] = {127, 0, 0, 1};

static kz_msgbox_id_t ip_mbox = MSGBOX_NULL;
static uint16 ip_id = 0;

void ip_init(void)
{
    ip_mbox = kz_getmsgbox();
}

void ip_handle(uint8 *packet)
{
    ASSERT(ip_mbox != MSGBOX_NULL, "no ip_mbox");
    if (kz_send(ip_mbox, 'h', (char *)packet, 0, NET_MBOXSIZE) == -1) {
        kz_kmfree(packet);
    }
}

void ip_send(uint8 *buf, uint16 size, uint8 protocol, const uint8 dst_IP[], const uint8 src_IP[])
{
    ip_header *ihdr = (ip_header *)(buf + sizeof(ethernet_header));

    ASSERT(ip_mbox != MSGBOX_NULL, "no ip_mbox");

    // 競合しない項目は呼出側で埋めてしまう
    ihdr->ip_version_length = 0x45;
    ihdr->ip_service_type = 0;
    ihdr->ip_total_length = sizeof(ip_header) + size;
    ihdr->ip_flags_fragment_offset = 0;
    ihdr->ip_time_to_live = 64;
    ihdr->ip_protocol = protocol;
    ihdr->ip_checksum = 0;
    memcpy(ihdr->ip_dst_IP, dst_IP, 4);
    if (src_IP) {
        memcpy(ihdr->ip_src_IP, src_IP, 4);
    } else {
        memset(ihdr->ip_src_IP, 0, 4); // NIが埋める。通常はこっち。
    }

    if (kz_send(ip_mbox, 's', (char *)buf, 0, NET_MBOXSIZE) == -1)
        kz_kmfree(buf);
}

void ip_send_data(uint8 *data, uint16 size, uint8 protocol, const uint8 dst_IP[])
{
    uint8 *buf = kz_kmalloc(sizeof(ethernet_header) + sizeof(ip_header) + size);
    memcpy(buf + sizeof(ethernet_header) + sizeof(ip_header), data, size);
    ip_send(buf, size, protocol, dst_IP, NULL);
}

static void handle_loopback(uint8 *buf)
{
    ip_header *ihdr = (ip_header *)(buf + sizeof(ethernet_header));

    ASSERT(is_loopback_address(ihdr->ip_dst_IP), "not loopback packet");

    if (memeq(ihdr->ip_dst_IP + 1, 0, 3) || memeq(ihdr->ip_dst_IP + 1, 255, 3)) {
        kz_kmfree(buf);
        return;
    }

    if (memeq(ihdr->ip_src_IP, 0, 4))
        memcpy(ihdr->ip_src_IP, loopback_address, 4);
    if (ihdr->ip_protocol == IP_PROTOCOL_UDP) {
        udp_header *uhdr = (udp_header *)(ihdr + 1);
        uhdr->udp_checksum = udp_checksum((uint8 *)uhdr, uhdr->udp_length,
                                          ihdr->ip_src_IP, ihdr->ip_dst_IP);
    }
    ihdr->ip_checksum = ip_checksum((uint8 *)ihdr, sizeof(ip_header));

    ip_handle(buf);
}

static void send_ip_packet(uint8 *buf)
{
    ip_header *ihdr = (ip_header *)(buf + sizeof(ethernet_header));

    ihdr->ip_id = ip_id++;
    // src_IPとチェックサムは以降で埋める

    if (is_loopback_address(ihdr->ip_dst_IP)) {
        handle_loopback(buf);
    } else {
        send_ni(buf);
    }
}

static int check_ihdr(ip_header *ihdr)
{
    int ihdr_len = (ihdr->ip_version_length & 0xf) << 2;

    return (ihdr->ip_version_length & 0xf0) == 0x40 &&      // IPv4
           !(ihdr->ip_flags_fragment_offset & ~0x4000) &&   // フラグメント不対応
           ihdr_len >= sizeof(*ihdr) &&
           ihdr_len <= ihdr->ip_total_length &&
           ip_checksum((uint8 *)ihdr, ihdr_len) == 0;
}

static void handle_option(ip_header *ihdr)
{
    // オプションは単に消す
    int ihdr_len = (ihdr->ip_version_length & 0xf) << 2;

    if (ihdr_len != (0x5 << 2)) {
        uint16 data_len = ihdr->ip_total_length - ihdr_len;

        ihdr->ip_version_length = (ihdr->ip_version_length & 0xf0) | 0x5;
        ihdr->ip_total_length = sizeof(ip_header) + data_len;

        memcpy(ihdr + 1, (uint8 *)ihdr + ihdr_len, data_len);
    }
}

void ip_main(int argc, char *argv)
{
    int cmd;
    char *p;

    while (1) {
        kz_recv(ip_mbox, &cmd, &p);

        if (cmd == 'h') { // 受信
            ip_header *ihdr = (ip_header *)(p + sizeof(ethernet_header));

            if (check_ihdr(ihdr)) {
                handle_option(ihdr);

                if (ihdr->ip_protocol == IP_PROTOCOL_ICMP) {
                    handle_icmp((uint8 *)p);
                    continue; // ここではpを解放しない

                } else if (ihdr->ip_protocol == IP_PROTOCOL_UDP) {
                    handle_udp((uint8 *)p);
                    continue;
                }
            }

            kz_kmfree(p);

        } else if (cmd == 's') { // 送信要求

            send_ip_packet((uint8 *)p);

        } else {
            ASSERT(0, "ip: Bad cmd");
        }
    }
}
